{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 预测部分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decoder(pred):\n",
    "    '''\n",
    "    输入：(1x7x7x30)---[x1,y1,x2,y2,置信度1,x1,y1,x2,y2,置信度1,....]\n",
    "    输出：\n",
    "    '''\n",
    "    grid_num = 14\n",
    "    cell_size = 1./grid_num\n",
    "    pred = pred.data\n",
    "    # 去掉batch_size维度--->(7x7x30)\n",
    "    pred = pred.squeeze(0)\n",
    "    # (7x7x1)\n",
    "    contain1 = pred[:, :, 4].unsqueeze(2)\n",
    "    # (7x7x1)\n",
    "    contain2 = pred[:, :, 9].unsqueeze(2)\n",
    "    # (7x7x2)\n",
    "    contain = torch.cat((contain1, contain2), 2)\n",
    "    # 把（7x7x2）里的数值替换为true, false\n",
    "    mask1 = contain > 0.1\n",
    "    mask2 = (contain == contain.max())\n",
    "    # 将mask1,mask2进行 或 操作\n",
    "    mask = (mask1 + mask2).gt(0)\n",
    "    \n",
    "    for i in range(grid_num):\n",
    "        for j in range(grid_num):\n",
    "            for b in range(2):\n",
    "                if mask[i, j, b] == 1:\n",
    "                    box = pred[i, j, b*5:b*5+4]\n",
    "                    contain_prob = torch.FloatTensor([pred[i, j, b*5+4]])\n",
    "                    ## 转换坐标\n",
    "                    xy = torch.FloatTensor([j, i]) * cell_size\n",
    "                    box[:2] = box[:2] * cell_size + xy\n",
    "                    \n",
    "                    max_prob, cls_index = torch.max(pred[i, j, 10:], 0)\n",
    "                    if float((contain_prob * max_prob)[0]) > 0.1:\n",
    "                        ## 添加"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
